import torch
import torch.nn as nn
import torch.nn.functional as F
import scipy.sparse as sp
import numpy as np
import torch.nn.init as init


class InputDropout(nn.Module):
    def __init__(self, keep_prob):
        super(InputDropout, self).__init__()
        self.p = keep_prob

    def forward(self, inputs):
        x = inputs.clone()
        if self.training:
            random_tensor = self.p + torch.rand((inputs.size(0),))
            dropout_mask = torch.floor(random_tensor).bool()
            x[~dropout_mask] = 0.
            return x / self.p
        else:
            return x


class StackGCNEncoder(nn.Module):
    def __init__(self, input_dim, output_dim, num_support,
                 dropout=0.,
                 use_bias=False, activation=F.relu):
        """对得到的每类评分使用级联的方式进行聚合

        Args:
        ----
            input_dim (int): 输入的特征维度
            output_dim (int): 输出的特征维度，需要output_dim % num_support = 0
            num_support (int): 评分的类别数，比如1~5分，值为5
            use_bias (bool, optional): 是否使用偏置. Defaults to False.
            activation (optional): 激活函数. Defaults to F.relu.
        """
        super(StackGCNEncoder, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.num_support = num_support
        self.dropout = dropout
        self.use_bias = use_bias
        self.activation = activation
        assert output_dim % num_support == 0
        self.weight = nn.Parameter(torch.Tensor(input_dim, output_dim))
        if self.use_bias:
            self.bias = nn.Parameter(torch.Tensor(output_dim, ))
        self.dropout = InputDropout(1 - dropout)
        self.reset_parameters()

    def reset_parameters(self):
        init.xavier_uniform_(self.weight)
        if self.use_bias:
            init.zeros_(self.bias)

    def forward(self, supports, inputs):
        """StackGCNEncoder计算逻辑

        Args:
            user_supports (list of torch.sparse.FloatTensor):
                归一化后每个评分等级对应的用户与商品邻接矩阵
            item_supports (list of torch.sparse.FloatTensor):
                归一化后每个评分等级对应的商品与用户邻接矩阵
            user_inputs (torch.Tensor): 用户特征的输入
            item_inputs (torch.Tensor): 商品特征的输入

        Returns:
            [torch.Tensor]: 用户的隐层特征
            [torch.Tensor]: 商品的隐层特征
        """
        assert len(supports) == self.num_support
        inputs = self.dropout(inputs)

        hidden = []
        weights = torch.split(self.weight, self.output_dim // self.num_support, dim=1)
        for i in range(self.num_support):
            tmp = torch.matmul(inputs, weights[i])
            tmp_hidden = torch.sparse.mm(supports[i][0], tmp)
            hidden.append(tmp_hidden)

        hidden = torch.cat(hidden, dim=1)

        outputs = self.activation(hidden)

        if self.use_bias:
            outputs += self.bias

        return outputs


class SumGCNEncoder(nn.Module):
    def __init__(self, input_dim, output_dim, num_support,
                 dropout=0.,
                 use_bias=False, activation=F.relu):
        """对得到的每类评分使用求和的方式进行聚合

        Args:
            input_dim (int): 输入的特征维度
            output_dim (int): 输出的特征维度，需要output_dim % num_support = 0
            num_support (int): 评分的类别数，比如1~5分，值为5
            use_bias (bool, optional): 是否使用偏置. Defaults to False.
            activation (optional): 激活函数. Defaults to F.relu.
        """
        super(SumGCNEncoder, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.num_support = num_support
        self.use_bias = use_bias
        self.activation = activation
        self.weight = nn.Parameter(torch.Tensor(
            input_dim, output_dim * num_support))
        if self.use_bias:
            self.bias = nn.Parameter(torch.Tensor(output_dim, ))
        self.dropout = InputDropout(1 - dropout)
        self.reset_parameters()

    def reset_parameters(self):
        init.xavier_uniform_(self.weight)
        if self.use_bias:
            init.zeros_(self.bias)

    def forward(self, supports, inputs):
        """SumGCNEncoder计算逻辑

        Args:
            user_supports (list of torch.sparse.FloatTensor):
                归一化后每个评分等级对应的用户与商品邻接矩阵
            item_supports (list of torch.sparse.FloatTensor):
                归一化后每个评分等级对应的商品与用户邻接矩阵
            user_inputs (torch.Tensor): 用户特征的输入
            item_inputs (torch.Tensor): 商品特征的输入

        Returns:
            [torch.Tensor]: 用户的隐层特征
            [torch.Tensor]: 商品的隐层特征
        """
        assert len(supports) == self.num_support
        inputs = self.dropout(inputs)

        hidden = []
        weights = torch.split(self.weight, self.output_dim, dim=1)
        for i in range(self.num_support):
            w = sum(weights[:(i + 1)])
            tmp = torch.matmul(inputs, w)
            tmp_hidden = torch.sparse.mm(supports[i], tmp)
            hidden.append(tmp_hidden)

        hidden = sum(hidden)
        outputs = self.activation(hidden)

        if self.use_bias:
            outputs += self.bias_user

        return outputs


class FullyConnected(nn.Module):
    def __init__(self, input_dim, output_dim, dropout=0.,
                 use_bias=False, activation=F.relu,
                 ):
        """非线性变换层

        Args:
        ----
            input_dim (int): 输入的特征维度
            output_dim (int): 输出的特征维度，需要output_dim % num_support = 0
            use_bias (bool, optional): 是否使用偏置. Defaults to False.
            activation (optional): 激活函数. Defaults to F.relu.
            share_weights (bool, optional): 用户和商品是否共享变换权值. Defaults to False.

        """
        super(FullyConnected, self).__init__()
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.use_bias = use_bias
        self.activation = activation
        self.weights = nn.Parameter(torch.Tensor(input_dim, output_dim))
        if use_bias:
            self.bias = nn.Parameter(torch.Tensor(output_dim))
        self.dropout = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self):

        init.xavier_uniform_(self.weights)
        if self.use_bias:
            init.normal_(self.bias, std=0.5)

    def forward(self, inputs):
        """前向传播

        Args:
            user_inputs (torch.Tensor): 输入的用户特征
            item_inputs (torch.Tensor): 输入的商品特征

        Returns:
            [torch.Tensor]: 输出的用户特征
            [torch.Tensor]: 输出的商品特征
        """
        x = self.dropout(inputs)
        x = torch.matmul(x, self.weights)

        outputs = self.activation(x)

        if self.use_bias:
            outputs += self.bias

        return outputs


class Decoder(nn.Module):
    def __init__(self, input_dim, num_weights, num_classes, dropout=0., activation=F.relu):
        """解码器

        Args:
        ----
            input_dim (int): 输入的特征维度
            num_weights (int): basis weight number
            num_classes (int): 总共的评分级别数，eg. 5
        """
        super(Decoder, self).__init__()
        self.input_dim = input_dim
        self.num_weights = num_weights
        self.num_classes = num_classes
        self.activation = activation

        self.weight = nn.ParameterList([nn.Parameter(torch.Tensor(input_dim, input_dim))
                                        for _ in range(num_weights)])
        self.weight_classifier = nn.Parameter(torch.Tensor(num_weights, num_classes))

        self.dropout = nn.Dropout(dropout)
        self.reset_parameters()

    def reset_parameters(self):
        for i in range(len(self.weight)):
            init.orthogonal_(self.weight[i], gain=1.1)
        init.xavier_uniform_(self.weight_classifier)

    def forward(self, inputs, source_indices, target_indices):
        """计算非归一化的分类输出

        Args:
            user_inputs (torch.Tensor): 用户的隐层特征
            item_inputs (torch.Tensor): 商品的隐层特征
            user_indices (torch.LongTensor):
                所有交互行为中用户的id索引，与对应的item_indices构成一条边,shape=(num_edges, )
            item_indices (torch.LongTensor):
                所有交互行为中商品的id索引，与对应的user_indices构成一条边,shape=(num_edges, )

        Returns:
            [torch.Tensor]: 未归一化的分类输出，shape=(num_edges, num_classes)
        """
        inputs = self.dropout(inputs)
        source_inputs = inputs[source_indices]
        target_inputs = inputs[target_indices]

        basis_outputs = []
        for i in range(self.num_weights):
            tmp = torch.matmul(source_inputs, self.weight[i])
            out = torch.sum(tmp * target_inputs, dim=1, keepdim=True)
            basis_outputs.append(out)

        basis_outputs = torch.cat(basis_outputs, dim=1)

        outputs = torch.matmul(basis_outputs, self.weight_classifier)
        outputs = self.activation(outputs)

        return outputs


class GraphMatrixCompletion(nn.Module):
    def __init__(self, input_dim, side_feat_dim,
                 gcn_hidden_dim, side_hidden_dim,
                 encode_hidden_dim,
                 num_support=26, num_classes=2, num_basis=3):
        super(GraphMatrixCompletion, self).__init__()
        self.encoder = StackGCNEncoder(input_dim, gcn_hidden_dim, num_support, dropout=0, use_bias=True)
        self.dense1 = FullyConnected(side_feat_dim, side_hidden_dim, dropout=0,
                                     use_bias=True)
        self.dense2 = FullyConnected(gcn_hidden_dim + side_hidden_dim, encode_hidden_dim,
                                     dropout=0, activation=lambda x: x)
        self.decoder = Decoder(encode_hidden_dim, num_basis, num_classes,
                               dropout=0, activation=lambda x: x)

    def forward(self, supports, inputs, side_inputs,
                source_edge_idx, target_edge_idx):
        gcn = self.encoder(supports, inputs)
        side_feat = self.dense1(side_inputs)

        feat = torch.cat((gcn, side_feat), dim=1)

        embed = self.dense2(feat)

        edge_logits = self.decoder(embed, source_edge_idx, target_edge_idx)

        return edge_logits

class GraphMatrixCompletion2(nn.Module):
    def __init__(self, input_dim, side_feat_dim,
                 gcn_hidden_dim, side_hidden_dim,
                 encode_hidden_dim,
                 num_support=26, num_classes=2, num_basis=3):
        super(GraphMatrixCompletion2, self).__init__()
        self.input_dim = input_dim
        self.gcn_hidden_dim = gcn_hidden_dim
        self.num_support = num_support
        self.encoder1 = StackGCNEncoder(input_dim, gcn_hidden_dim, num_support, dropout=0, use_bias=True)
        self.encoder2 = StackGCNEncoder(gcn_hidden_dim, gcn_hidden_dim, num_support, dropout=0, use_bias=True)

        self.dense1 = FullyConnected(side_feat_dim, side_hidden_dim, dropout=0,
                                     use_bias=True)
        self.dense2 = FullyConnected(gcn_hidden_dim + side_hidden_dim, encode_hidden_dim,
                                     dropout=0, activation=lambda x: x)
        self.decoder = Decoder(encode_hidden_dim, num_basis, num_classes,
                               dropout=0, activation=lambda x: x)

    def forward(self, supports, inputs, side_inputs,
                source_edge_idx, target_edge_idx):
        gcn1 = self.encoder1(supports, inputs)
        gcn2 = self.encoder2(supports, gcn1)

        side_feat = self.dense1(side_inputs)

        feat = torch.cat((gcn2, side_feat), dim=1)

        embed = self.dense2(feat)

        edge_logits = self.decoder(embed, source_edge_idx, target_edge_idx)

        return edge_logits


class GraphMatrixCompletion3(nn.Module):
    def __init__(self, input_dim, side_feat_dim,
                 gcn_hidden_dim, side_hidden_dim,
                 encode_hidden_dim,
                 num_support=26, num_classes=2, num_basis=3):
        super(GraphMatrixCompletion3, self).__init__()
        self.input_dim = input_dim
        self.gcn_hidden_dim = gcn_hidden_dim
        self.num_support = num_support
        self.encoder1 = StackGCNEncoder(input_dim, gcn_hidden_dim//2, num_support, dropout=0, use_bias=True)
        self.encoder2 = StackGCNEncoder(gcn_hidden_dim//2, gcn_hidden_dim, num_support, dropout=0, use_bias=True)
        self.encoder3 = StackGCNEncoder(gcn_hidden_dim, gcn_hidden_dim, num_support, dropout=0, use_bias=True)

        self.dense1 = FullyConnected(side_feat_dim, side_hidden_dim, dropout=0,
                                     use_bias=True)
        self.dense2 = FullyConnected(gcn_hidden_dim + side_hidden_dim, encode_hidden_dim,
                                     dropout=0, activation=lambda x: x)
        self.decoder = Decoder(encode_hidden_dim, num_basis, num_classes,
                               dropout=0, activation=lambda x: x)

    def forward(self, supports, inputs, side_inputs,
                source_edge_idx, target_edge_idx):
        gcn1 = self.encoder1(supports, inputs)
        gcn2 = self.encoder2(supports, gcn1)
        gcn3 = self.encoder3(supports, gcn2)

        side_feat = self.dense1(side_inputs)

        feat = torch.cat((gcn3, side_feat), dim=1)

        embed = self.dense2(feat)

        edge_logits = self.decoder(embed, source_edge_idx, target_edge_idx)

        return edge_logits

